from pymilvus import (
    connections,
    utility,
    CollectionSchema,
    Collection,
)
import logging

logger = logging.getLogger(__name__)


class MilvusClient:

    def __init__(self,host: str = "localhost", port: str = "19530", collection_name: str = "chat_records", user:str="root",pwd:str="szsk8888",fields = []):
        self.user = user
        self.pwd = pwd
        self.host = host
        self.port = port
        self.collection_name = collection_name
        self.collection = None
        self._connected = False
        self.collection_fields=fields
        self.ensure_collection_ready()


    def connect(self):
        """安全连接方法"""
        if self._connected:
            try:
                handler = connections._get_connection_handler('default')
                if handler._host == self.host and handler._port == self.port:
                    return
                connections.disconnect('default')
            except:
                pass

        try:
            connections.connect(
                alias='default',
                host=self.host,
                port=self.port,
                user=self.user,
                password=self.pwd,
                secure=False
                # 可添加其他参数如: user, password, secure 等
            )
            self._connected = True
            logger.info(f"Milvus connected to {self.host}:{self.port}")
        except Exception as e:
            logger.error(f"Connection failed: {str(e)}")
            raise

    def ensure_collection_ready(self):
        """
        确保集合已准备就绪
        新增方法替代原来的 _ensure_collection_ready
        """
        if not self._connected:
            self.connect()

        if self.collection is None:
            self.init_collection()

        #if not self.collection.has_index():
            #self.create_index()

    def init_collection(self, dimension=384):
        """初始化集合"""
        self.connect()

        if utility.has_collection(self.collection_name):
            self.collection = Collection(self.collection_name)
            logger.info(f"Loaded existing collection: {self.collection_name}")
        else:
            self.create_collection(dimension)

        #if not self.collection.has_index():
            #self.create_index()

    def create_collection(self, dimension: int):
        """创建新集合"""
        schema = CollectionSchema(self.collection_fields, description="Enhanced chat records with vector embeddings")
        self.collection = Collection(self.collection_name, schema)
        logger.info(f"已创建新集合: {self.collection_name}")

        self.create_index()

    def create_index(self):
        """创建向量索引"""
        index_params = {
            "metric_type": "L2",
            "index_type": "IVF_FLAT",
            "params": {"nlist": 128}
        }

        self.collection.create_index(
            field_name="embedding",
            index_params=index_params
        )
        logger.info("已创建向量索引")




    def load_collection(self):
        if self.collection:
            self.collection.load()
